import datetime
import logging
import time
import traceback
import kafka

from coin.base.datetime_util import to_timestamp_int
from coin.base.timestamp import get_timestamp
from coin.proto.coin_feed_pb2 import SubscriptionLog
from coin.tool.strat_monitor.util.monitor_util import _filter_ticker_topic
from coin.util.queue.tools.kafka_archive import run_from_kafka_archive
from xunkemgmt_client.support.accounting.database.db_output import DBOutputTicker
from xunkemgmt_client.support.accounting.database.db_querier import MarketInfoQuerier


def _validate_sub_request(account_request):
  assert account_request.HasField('market_type')
  assert account_request.HasField('exchange')
  assert account_request.HasField('api_version')


class KafkaMarketInfoQuerier(MarketInfoQuerier):
  def __init__(self, kafka_config):
    self._kafka_config = kafka_config
    self._topics = kafka.KafkaConsumer(
        bootstrap_servers=self._kafka_config.kafka_servers).topics()
    self._ticker = {}
    self._sub_key_func = lambda x: (x.market_type, x.exchange, x.api_version)

  def _on_log(self, partition, record):
    pb = SubscriptionLog()
    pb.ParseFromString(record.value)
    assert pb.HasField('request')
    _validate_sub_request(pb.request)
    sub_key = self._sub_key_func(pb.request)
    if pb.type == SubscriptionLog.TICKER:
      assert pb.HasField('ticker')
      if len(pb.ticker.each_ticker) > 0:
        self._ticker[sub_key] = pb.ticker
    else:
      pass

  def query_ticker_history(self, request_list, query_timestamp=None):
    self._ticker.clear()
    topics = _filter_ticker_topic(self._topics, request_list)
    timestamp_to = query_timestamp or get_timestamp()
    timestamp_from = timestamp_to - 10 * 60 * 10**9
    run_from_kafka_archive(topics=topics,
                           kafka_config=self._kafka_config,
                           timestamp_from=timestamp_from,
                           timestamp_to=timestamp_to,
                           callback=self._on_log)
    result_list = []

    for request in request_list:
      try:
        assert request.exchange_api_id is not None
        ticker = self._ticker[self._sub_key_func(request)]
        ticker_result = DBOutputTicker(
            exchange_api_id=request.exchange_api_id,
            query_timestamp=time.time(),
            ticker=ticker,
            ticker_hid=None)
        result_list.append(ticker_result)
      except Exception as e:
        logging.error(f'Fail to query ticker {request}')
        logging.error(traceback.format_exc())
        result_list.append(None)
        continue
    return result_list


if __name__ == "__main__":
  from coin.proto.coin_request_pb2 import AccountRequestProto
  from coin.util.queue.config import KafkaConfig

  kafka_config = KafkaConfig.from_cmd_config(
      '../../coin_deploy/support_monitor/config/kafka_aws_config_2.json')
  querier = KafkaMarketInfoQuerier(kafka_config)
  request_list = [AccountRequestProto(
      market_type='Spot', exchange='Binance', api_version='v1')]
  query_timestamp = None
  query_timestamp = to_timestamp_int(datetime.datetime(2023,3,14))
  print(querier.query_ticker_history(request_list, query_timestamp=query_timestamp))
